GraphSage Embeddings Guide

This guide will show how to create graph sage embeddings. GraphSage training and models does not work in the current neo4j training and link prediction pipeline.

Project A Graph

Project a graph for the pubmed articles. Include drug, disease, gene_protein foundational nodes. Include only the `has_extraction` pubmed article relationship to the foundational graph nodes. Include the `model2vec_embeddings` to be merged with the graph sage embeddings.

call gds.graph.project(
  'pubmed_graph',
  [
      'disease', 
      'drug',
      'gene_protein',
      'pubmed_document'
  ],
  {
    has_extraction: {orientation: 'UNDIRECTED', type:'*'},
    disease_disease: {orientation: 'UNDIRECTED', type:'*'},
    disease_protein: {orientation: 'UNDIRECTED', type:'*'},
    drug_protein: {orientation: 'UNDIRECTED', type:'*'},
    drug_drug: {orientation: 'UNDIRECTED', type:'*'},
    drug_effect: {orientation: 'UNDIRECTED', type:'*'},
    protein_protein: {orientation: 'UNDIRECTED', type:'*'},
    indication: {orientation: 'UNDIRECTED', type:'*'},
    contraindication: {orientation: 'UNDIRECTED', type:'*'}
  },
  {
    nodeProperties: ['model2vec_embeddings', 'node_idx', 'label_one_hot_encoding']
  }
)
yield
  graphName AS graph, nodeProjection, nodeCount AS nodes, relationshipCount AS rels
return
  graph, nodeProjection, nodes, rels;

Project a pubmed graph

Sample Graph For Training

Reduce the size of the training graph using https://neo4j.com/docs/graph-data-science/current/management-ops/graph-creation/sampling/rwr/[random walk with restarts]. We can use random walk to train on a subset of the total graph because graph sage is "inductive" meaning we can learn embeddings for new unseen nodes.

call gds.graph.sample.rwr('pubmed_sampled_graph', 'pubmed_graph', { samplingRatio: 0.35 })
yield nodeCount, relationshipCount
return nodeCount, relationshipCount;

Reduce graph size by sampling with random walks

Train Graph Embeddings

Train graph embeddings using GraphSage. Merges graph embeddings with the sentence embeddings in field `model2vec_embeddings`.

call gds.beta.graphSage.train(
  'pubmed_sampled_graph',
  {
      modelName: 'pubmed_graphsage_model',
      featureProperties: ['model2vec_embeddings', 'label_one_hot_encoding'],
      embeddingDimension: 512,
      aggregator: 'mean',
      epochs: 3,
      maxIterations: 15,
      learningRate: .075,
      batchSize: 128,
      searchDepth: 5,
      penaltyL2: .001,
      sampleSizes: [32, 12],
      activationFunction: 'relu',
      projectedFeatureDimension: 256
  }
)
yield trainMillis, configuration, modelInfo
return trainMillis, configuration, modelInfo;

Train graph embeddings using GraphSage

Write Graph Embeddings To A New Field

Write graph embeddings to new field.

call gds.beta.graphSage.mutate(
  'pubmed_graph',
  {
    mutateProperty: 'graphsage_embedding',
    modelName: 'pubmed_graphsage_model'
  }
) yield
  nodeCount,
  nodePropertiesWritten;

Write graph embeddings

Use Graph Embeddings To Do KNN Node Classification

Use graph embeddings to do KNN node classification. The classification will connect similar PubMed articles.

call gds.knn.filtered.write('pubmed_graph', {
  nodeProperties: [
    {
      graphsage_embedding: 'COSINE'
    }
  ],
  nodeLabels: ['pubmed_document'],
  topK: 6,
  sampleRate: 1.0,
  deltaThreshold: 0.001,
  similarityCutoff: .85,
  randomSeed: 75,
  concurrency: 1,
  writeProperty: 'graphsage_score',
  writeRelationshipType: 'similar_article'
})
yield similarityDistribution
return similarityDistribution.mean as meanSimilarity;

Use graph embeddings to do KNN node classification

Clean Up

Delete graph, models and pipeline.

// delete article knn links
match (d1:pubmed_document)-[r:similar_article]-(d2:pubmed_document)
delete r;

// drop model
call gds.model.drop('pubmed_graphsage_model');

// drop graphs
call gds.graph.drop('pubmed_graph');
call gds.graph.drop('pubmed_sampled_graph');

Clean up